from typing import Optional, Any, Tuple, List
from fastapi import FastAPI
from fastapi.responses import ORJSONResponse as JSONResponse
from fastapi import status
import math


class APIResponse(JSONResponse):
    def __init__(self, code: int = 100, msg: Optional[str] = '成功', status_code: int = status.HTTP_200_OK,
                 **kwargs) -> None:
        self.data = {
            'code': code,
            'msg': msg
        }
        self.data.update(kwargs)
        super().__init__(content=self.data, status_code=status_code)


class PaginationResponse(JSONResponse):
    def __init__(self,
                 data: List[Any] = [],
                 msg: Optional[str] = "查询成功",
                 page: int = 1,
                 page_size: int = 10,
                 code: int = 100,
                 status_code: int = status.HTTP_200_OK
                 ) -> None:
        if page < 1 or page_size < 1:
            raise

        total, results = self.get_paginate_response(data, page, page_size)
        has_next = True if math.ceil(total / page_size) > page else False
        self.data = {
            "code": code,
            "msg": msg,
            "total": total,
            "page": page,
            "page_size": page_size,
            "has_next": has_next,
            "results": results,
        }
        super().__init__(content=self.data, status_code=status_code)

    @staticmethod
    def get_paginate_response(data, page, page_size) -> Tuple[int, List[Any]]:
        start = (page - 1) * page_size
        end = page * page_size
        paginated_data = data[start:end]
        return len(data), paginated_data
